#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import logging
import os
import sys
import subprocess
from datetime import datetime
import time

# 添加项目路径到系统路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger("ClusterScheduler")

class ClusterScheduler:
    def __init__(self, config_path="config.json"):
        """初始化定时任务调度器"""
        self.config_path = config_path
        self.config = self._load_config()
        
    def _load_config(self):
        """加载配置文件"""
        try:
            with open(self.config_path, 'r') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}")
            return {}
    
    def setup_cron_job(self):
        """设置cron定时任务"""
        try:
            # 获取当前脚本的绝对路径
            script_path = os.path.abspath(__file__)
            cluster_inspection_path = os.path.join(os.path.dirname(script_path), "cluster-inspection.py")
            python_path = sys.executable
            
            # 构造cron命令
            cron_command = f"{python_path} {cluster_inspection_path}"
            
            # 获取定时任务表达式
            schedule = self.config.get("schedule", "0 9 * * *")
            
            # 构造cron条目
            cron_entry = f"{schedule} {cron_command}"
            
            # 获取当前用户的crontab
            try:
                current_crontab = subprocess.check_output(["crontab", "-l"], stderr=subprocess.STDOUT).decode()
            except subprocess.CalledProcessError:
                current_crontab = ""
            
            # 检查是否已存在相同的任务
            if cron_command in current_crontab:
                # 如果存在，先移除旧的任务
                lines = current_crontab.strip().split('\n')
                new_lines = []
                for line in lines:
                    if line.strip() and cron_command not in line:
                        new_lines.append(line)
                new_lines.append(cron_entry)
                updated_crontab = '\n'.join(new_lines) + '\n'
            else:
                # 如果不存在，直接添加
                updated_crontab = current_crontab + cron_entry + '\n'
            
            # 更新crontab
            process = subprocess.Popen(["crontab", "-"], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            stdout, stderr = process.communicate(input=updated_crontab.encode())
            
            if process.returncode == 0:
                logger.info(f"定时任务设置成功: {cron_entry}")
                return True
            else:
                logger.error(f"设置定时任务失败: {stderr.decode()}")
                return False
                
        except Exception as e:
            logger.error(f"设置定时任务时发生错误: {e}")
            return False
    
    def remove_cron_job(self):
        """移除cron定时任务"""
        try:
            # 获取当前脚本的绝对路径
            script_path = os.path.abspath(__file__)
            cluster_inspection_path = os.path.join(os.path.dirname(script_path), "cluster-inspection.py")
            python_path = sys.executable
            
            # 构造cron命令
            cron_command = f"{python_path} {cluster_inspection_path}"
            
            # 获取当前用户的crontab
            try:
                current_crontab = subprocess.check_output(["crontab", "-l"], stderr=subprocess.STDOUT).decode()
            except subprocess.CalledProcessError:
                logger.info("当前没有设置任何定时任务")
                return True
            
            # 移除包含cron_command的行
            lines = current_crontab.strip().split('\n')
            new_lines = []
            removed = False
            for line in lines:
                if line.strip() and cron_command not in line:
                    new_lines.append(line)
                elif cron_command in line:
                    removed = True
                    
            if removed:
                updated_crontab = '\n'.join(new_lines) + '\n' if new_lines else ""
                # 更新crontab
                process = subprocess.Popen(["crontab", "-"], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                stdout, stderr = process.communicate(input=updated_crontab.encode())
                
                if process.returncode == 0:
                    logger.info("定时任务移除成功")
                    return True
                else:
                    logger.error(f"移除定时任务失败: {stderr.decode()}")
                    return False
            else:
                logger.info("未找到相关的定时任务")
                return True
                
        except Exception as e:
            logger.error(f"移除定时任务时发生错误: {e}")
            return False
    
    def list_cron_jobs(self):
        """列出当前用户的cron定时任务"""
        try:
            current_crontab = subprocess.check_output(["crontab", "-l"], stderr=subprocess.STDOUT).decode()
            print("当前用户的定时任务:")
            print(current_crontab)
            return current_crontab
        except subprocess.CalledProcessError as e:
            if "no crontab" in str(e.output):
                print("当前用户没有设置任何定时任务")
            else:
                print(f"获取定时任务列表失败: {e.output.decode()}")
            return ""
        except Exception as e:
            logger.error(f"获取定时任务列表时发生错误: {e}")
            return ""
    
    def run_once(self):
        """立即运行一次巡检任务"""
        try:
            # 获取当前脚本的绝对路径
            script_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cluster-inspection.py")
            python_path = sys.executable
            
            logger.info("开始执行集群巡检任务...")
            result = subprocess.run([python_path, script_path], capture_output=True, text=True)
            
            if result.returncode == 0:
                logger.info("集群巡检任务执行成功")
                print(result.stdout)
                return True
            else:
                logger.error(f"集群巡检任务执行失败: {result.stderr}")
                return False
                
        except Exception as e:
            logger.error(f"执行巡检任务时发生错误: {e}")
            return False

def main():
    """主函数"""
    if len(sys.argv) < 2:
        print("使用方法:")
        print("  python scheduler.py setup    - 设置定时任务")
        print("  python scheduler.py remove   - 移除定时任务")
        print("  python scheduler.py list     - 列出定时任务")
        print("  python scheduler.py run      - 立即运行一次")
        return
    
    scheduler = ClusterScheduler()
    
    command = sys.argv[1]
    if command == "setup":
        scheduler.setup_cron_job()
    elif command == "remove":
        scheduler.remove_cron_job()
    elif command == "list":
        scheduler.list_cron_jobs()
    elif command == "run":
        scheduler.run_once()
    else:
        print("未知命令，请使用 setup, remove, list 或 run")

if __name__ == "__main__":
    main()